from app.util import Config

from langchain_openai import ChatOpenAI
from langchain.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_community.chat_message_histories import ChatMessageHistory
from langchain_core.messages import AIMessage
from langchain_core.runnables.history import RunnableWithMessageHistory
import tomli


class AiChater:
    def __init__(self, config: Config):
        with open(config.model_config, 'rb') as file: 
            model_config = tomli.load(file)

        model = ChatOpenAI(
            api_key=model_config['llm']['api_key'],
            model=model_config['llm']['model'],
            openai_api_base=model_config['llm']['base_url'],
        )

        prompt_template = ChatPromptTemplate.from_messages([
            ('system', 'You are a chatbot'),
            MessagesPlaceholder(variable_name='chat_history'),
            ('human', '{new_message}'),
        ])

        chain = prompt_template | model
        
        self.chat_history = ChatMessageHistory()

        self.chain_with_memory = RunnableWithMessageHistory(
            chain,
            lambda x: self.chat_history,
            input_messages_key='new_message',
            history_messages_key='chat_history',
        )

        summarize_prompt = ChatPromptTemplate.from_messages([
            MessagesPlaceholder(variable_name='chat_history'),
            ('human', '将上述聊天信息提炼为一条摘要信息，尽可能包含具体细节。')
        ])
        self.summarize_chain = summarize_prompt | model

    def invoke(self, user_input: str):
        self._summarize_message()
        response = self.chain_with_memory.invoke(
            {"new_message": user_input},
            config={"configurable": {"session_id": "chat-gui"}}
        )
        return response.content

    def _summarize_message(self):
        stored_messages = self.chat_history.messages
        if len(stored_messages) > 6:
            result = self.summarize_chain.invoke({'chat_history': stored_messages})
            summary_content = result.content
            self.chat_history.clear()
            self.chat_history.add_messages([AIMessage(content=summary_content)])

